perf: Cache Edwards public key for XEdDSA signatures (-27% signing, with lazy key gen) - #240
Merged
Merged
Conversation
📝 WalkthroughWalkthroughCaches Edwards public key and sign bit inside private-key representations to avoid repeated scalar multiplications; updates signing, key generation, and serialization to use cached values. Replaces several hash finalizations with zero-allocation helpers, removes one SessionState accessor, and converts multiple ownership moves into explicit clones. Changes
Sequence Diagram(s)(Skipped) Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Bencher Report
jlucaso1
force-pushed
the
perf-experimental-cache-curve
branch
from
January 24, 2026 12:38
289041b to
799fd6d
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@wacore/libsignal/src/protocol/state/session.rs`:
- Around line 227-247: The fast-path in get_receiver_chain_index currently
compares raw sender bytes to stored chain.sender_ratchet_key bytes without
validating format, which can hide corrupted keys and later cause panics in
set_message_keys / set_receiver_chain_key; add a cheap validation gate before
the byte-comparison that verifies the stored sender_ratchet_key has the expected
serialized layout (e.g., correct type/version byte and length matching
PublicKey's serialized size or the constant used by PublicKey::serialize),
returning Err(InvalidSessionError("invalid receiver chain ratchet key")) on
failure; keep the direct byte comparison afterwards to preserve the optimization
while restoring the original validation semantics for sender_ratchet_key.
🧹 Nitpick comments (1)
wacore/libsignal/src/core/curve/curve25519.rs (1)
70-83: Masksign_bitto avoid invalid signatures from bad inputs.If a caller passes a value other than
0x00or0x80, the signature encoding can be corrupted. A small mask keeps the API safe.♻️ Suggested hardening
- PrivateKey { - secret, - ed_public_key, - sign_bit, - } + PrivateKey { + secret, + ed_public_key, + sign_bit: sign_bit & 0b1000_0000_u8, + }
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Fix all issues with AI agents
In `@wacore/libsignal/src/core/curve.rs`:
- Around line 366-377: The generate function wastes work by creating temp via
curve25519::PrivateKey::new(csprng) which computes the Edwards cache, then
discarding it when building PrivateKeyData::DjbPrivateKey with edwards_cache:
OnceLock::new(); change generate to preserve temp's computed edwards cache into
the new PrivateKeyData (e.g., initialize the OnceLock with the computed cache or
transfer the cached representation from temp into edwards_cache) so the
precomputed cache from temp is reused; locate symbols generate,
curve25519::PrivateKey::new, temp, PrivateKeyData::DjbPrivateKey, and
edwards_cache/OnceLock to implement the transfer.
- Around line 216-227: PrivateKeyData cannot derive Clone because
OnceLock<EdwardsCacheData> isn't Clone; replace the edwards_cache field with a
cloneable sync wrapper such as Arc<Mutex<Option<EdwardsCacheData>>> (e.g.,
edwards_cache: Arc<Mutex<Option<EdwardsCacheData>>>), update any initialization
to Arc::new(Mutex::new(None)), and update code paths that lazily compute or read
the Edwards cache (where edwards_cache is accessed during signing) to lock the
mutex, check/set the Option, and return the cached value; this makes
PrivateKeyData::DjbPrivateKey cloneable while preserving the lazy-cache
semantics.
In `@wacore/libsignal/src/core/curve/curve25519.rs`:
- Around line 99-110: from_bytes_without_cache currently constructs a PrivateKey
with dummy ed_public_key and sign_bit=0 which can silently produce invalid
signatures if calculate_signature is called; update this by choosing a clear
sentinel (e.g., an out-of-range sign_bit value or a reserved ed_public_key
pattern) when constructing the PrivateKey in from_bytes_without_cache and then
make calculate_signature check that sentinel and return an explicit error or
panic instead of producing a signature; update the doc comment on
from_bytes_without_cache to warn about non-signing use and reference the
sentinel so future maintainers see the contract (symbols:
from_bytes_without_cache, PrivateKey, ed_public_key, sign_bit,
calculate_signature).
🧹 Nitpick comments (1)
wacore/libsignal/src/core/curve/curve25519.rs (1)
70-85: Consider adding debug assertion for cache consistency.The
from_bytes_with_cacheconstructor trusts that the provideded_public_keyandsign_bitmatch theprivate_key. While this is intentional for performance, adebug_assert!could help catch mismatches during development/testing without affecting release builds.💡 Optional: Add debug assertion
pub fn from_bytes_with_cache( private_key: [u8; PRIVATE_KEY_LENGTH], ed_public_key: CompressedEdwardsY, sign_bit: u8, ) -> Self { let secret = StaticSecret::from(scalar::clamp_integer(private_key)); + #[cfg(debug_assertions)] + { + let (expected_ed, expected_sign) = Self::compute_ed_public_key(&secret); + debug_assert_eq!(ed_public_key, expected_ed, "Edwards public key mismatch"); + debug_assert_eq!(sign_bit & 0b1000_0000_u8, expected_sign, "Sign bit mismatch"); + } PrivateKey { secret, ed_public_key, // Mask to ensure only valid sign bit values (0x00 or 0x80) sign_bit: sign_bit & 0b1000_0000_u8, } }
Comment on lines
+216
to
+227
| /// Stores the private key bytes with lazy-initialized cached values for XEdDSA signing. | ||
| /// The Edwards public key is computed on first signature, not at key creation. | ||
| /// This keeps key generation fast while subsequent signatures benefit from caching. | ||
| #[derive(Debug, Clone)] | ||
| enum PrivateKeyData { | ||
| DjbPrivateKey([u8; curve25519::PRIVATE_KEY_LENGTH]), | ||
| DjbPrivateKey { | ||
| /// The raw 32-byte private key | ||
| key: [u8; curve25519::PRIVATE_KEY_LENGTH], | ||
| /// Lazily-initialized Edwards cache (computed on first signature) | ||
| edwards_cache: OnceLock<EdwardsCacheData>, | ||
| }, | ||
| } |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, let's examine the actual code at lines 216-227 in curve.rs
sed -n '200,240p' wacore/libsignal/src/core/curve.rs | cat -nRepository: jlucaso1/whatsapp-rust
Length of output: 1604
🏁 Script executed:
# Check the OnceLock import and usage
rg -n "use.*OnceLock|from.*OnceLock" wacore/libsignal/src/core/curve.rsRepository: jlucaso1/whatsapp-rust
Length of output: 94
🏁 Script executed:
# Check for any Clone implementations for PrivateKeyData
rg -n "impl.*Clone.*PrivateKeyData|impl Clone for PrivateKeyData" wacore/libsignal/src/core/curve.rsRepository: jlucaso1/whatsapp-rust
Length of output: 48
🏁 Script executed:
# Check the full PrivateKeyData enum definition
rg -B5 -A20 "enum PrivateKeyData" wacore/libsignal/src/core/curve.rs | head -40Repository: jlucaso1/whatsapp-rust
Length of output: 936
🏁 Script executed:
# Check Cargo.toml for Rust version requirements
rg -n "rust-version|edition" wacore/Cargo.tomlRepository: jlucaso1/whatsapp-rust
Length of output: 84
🏁 Script executed:
# Check if the code compiles by looking for build logs or test results
# First, check if there's a way to verify compilation
cd wacore && cargo check 2>&1 | head -100Repository: jlucaso1/whatsapp-rust
Length of output: 175
🏁 Script executed:
# Also check if there are any feature flags or conditional attributes that might affect OnceLock
rg -n "feature|cfg" wacore/libsignal/src/core/curve.rs | head -20Repository: jlucaso1/whatsapp-rust
Length of output: 48
OnceLock does not implement Clone, causing a compilation error.
PrivateKeyData cannot derive Clone because its OnceLock<EdwardsCacheData> field does not implement Clone. This will fail to compile with "the trait bound OnceLock<EdwardsCacheData>: Clone is not satisfied". Either remove the Clone derive or replace OnceLock with a synchronization primitive that implements Clone (e.g., Arc<Mutex<Option<EdwardsCacheData>>> or parking_lot::Once).
🤖 Prompt for AI Agents
In `@wacore/libsignal/src/core/curve.rs` around lines 216 - 227, PrivateKeyData
cannot derive Clone because OnceLock<EdwardsCacheData> isn't Clone; replace the
edwards_cache field with a cloneable sync wrapper such as
Arc<Mutex<Option<EdwardsCacheData>>> (e.g., edwards_cache:
Arc<Mutex<Option<EdwardsCacheData>>>), update any initialization to
Arc::new(Mutex::new(None)), and update code paths that lazily compute or read
the Edwards cache (where edwards_cache is accessed during signing) to lock the
mutex, check/set the Option, and return the cached value; this makes
PrivateKeyData::DjbPrivateKey cloneable while preserving the lazy-cache
semantics.
Comment on lines
+99
to
110
| /// Creates a PrivateKey from raw bytes WITHOUT computing the Edwards cache. | ||
| /// Use this for operations that don't need signatures (e.g., key agreement, public key derivation). | ||
| #[inline] | ||
| pub fn from_bytes_without_cache(private_key: [u8; PRIVATE_KEY_LENGTH]) -> Self { | ||
| let secret = StaticSecret::from(scalar::clamp_integer(private_key)); | ||
| // Use dummy values - these should never be accessed for non-signature operations | ||
| PrivateKey { | ||
| secret, | ||
| ed_public_key: CompressedEdwardsY::default(), | ||
| sign_bit: 0, | ||
| } | ||
| } |
There was a problem hiding this comment.
Silent failure risk if used for signing.
If a PrivateKey created via from_bytes_without_cache is accidentally used for calculate_signature, it will produce invalid signatures using dummy Edwards data. Consider adding a sentinel value or documentation warning.
💡 Suggested improvement
One option is to use a clearly invalid sentinel value that would cause signature verification to fail obviously, or add a comment at the calculate_signature call site. At minimum, strengthen the doc comment:
- /// Creates a PrivateKey from raw bytes WITHOUT computing the Edwards cache.
- /// Use this for operations that don't need signatures (e.g., key agreement, public key derivation).
+ /// Creates a PrivateKey from raw bytes WITHOUT computing the Edwards cache.
+ ///
+ /// # Warning
+ /// Keys created with this function MUST NOT be used for `calculate_signature`.
+ /// Doing so will produce invalid signatures. Use this only for key agreement
+ /// and public key derivation operations.🤖 Prompt for AI Agents
In `@wacore/libsignal/src/core/curve/curve25519.rs` around lines 99 - 110,
from_bytes_without_cache currently constructs a PrivateKey with dummy
ed_public_key and sign_bit=0 which can silently produce invalid signatures if
calculate_signature is called; update this by choosing a clear sentinel (e.g.,
an out-of-range sign_bit value or a reserved ed_public_key pattern) when
constructing the PrivateKey in from_bytes_without_cache and then make
calculate_signature check that sentinel and return an explicit error or panic
instead of producing a signature; update the doc comment on
from_bytes_without_cache to warn about non-signing use and reference the
sentinel so future maintainers see the contract (symbols:
from_bytes_without_cache, PrivateKey, ed_public_key, sign_bit,
calculate_signature).
Optimizations: 1. get_receiver_chain_index(): Compare serialized bytes directly instead of deserializing PublicKey on each iteration. Avoids allocation and validation overhead per chain. 2. Remove unused get_receiver_chain() method that was cloning protobuf Chain structs unnecessarily. 3. Use finalize_sha256_array() in upload.rs and messages.rs for zero-allocation hash finalization instead of finalize() + conversion.
- Mask sign_bit in from_bytes_with_cache() to prevent invalid signatures from bad inputs (only 0x00 or 0x80 are valid) - Add format validation in get_receiver_chain_index() before byte comparison to catch corrupted keys early instead of causing panics later - Make KeyType::value() public to avoid duplicating the 0x05 constant - Use KeyType::Djb.value() in session.rs validation for DRY
Add new_without_cache() function to curve25519::PrivateKey that generates keys without computing the expensive Edwards public key.
jlucaso1
force-pushed
the
perf-experimental-cache-curve
branch
from
January 24, 2026 14:02
536d03d to
8fd2d31
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@wacore/libsignal/src/core/curve.rs`:
- Around line 366-378: The comment is correct: KeyPair::generate currently calls
curve25519::PrivateKey::new_without_cache expecting no Edwards computation, but
new_without_cache still calls compute_ed_public_key; update the implementation
so it truly defers expensive scalar multiplication (or change the call site).
Specifically, modify curve25519::PrivateKey::new_without_cache to mirror
from_bytes_without_cache by NOT calling compute_ed_public_key and instead
initializing the edwards cache as an empty OnceLock/dummy placeholder, or
alternatively change KeyPair::generate to call from_bytes_without_cache (or
another clearly-named factory) so the API behavior matches the “lazy” claim
(referencing PrivateKey::new_without_cache,
PrivateKey::from_bytes_without_cache, and KeyPair::generate).
🧹 Nitpick comments (1)
wacore/libsignal/src/protocol/state/session.rs (1)
636-643: Avoidexpectin the promotion path.Even with a validated index, a panic here isn’t ideal; return a structured error instead to keep invariants defensive.
[warned_recommended_refactor? no] → use suggested tag.As per coding guidelines, avoid `unwrap`/`expect` outside tests.♻️ Suggested fix
- let state = self - .take_previous_session(index) - .expect("index was just validated"); + let state = self + .take_previous_session(index) + .ok_or(InvalidSessionError("missing previous session"))?; self.promote_state(state);
Comment on lines
366
to
+378
| pub fn generate<R: Rng + CryptoRng>(csprng: &mut R) -> Self { | ||
| let private_key = curve25519::PrivateKey::new(csprng); | ||
|
|
||
| let public_key = PublicKey::from(PublicKeyData::DjbPublicKey( | ||
| private_key.derive_public_key_bytes(), | ||
| )); | ||
| let private_key = PrivateKey::from(PrivateKeyData::DjbPrivateKey( | ||
| private_key.private_key_bytes(), | ||
| )); | ||
| // Generate key WITHOUT computing Edwards cache (lazy initialization). | ||
| // The Edwards point computation is deferred until first signature. | ||
| let temp = curve25519::PrivateKey::new_without_cache(csprng); | ||
| let key = temp.private_key_bytes(); | ||
|
|
||
| let public_key = | ||
| PublicKey::from(PublicKeyData::DjbPublicKey(temp.derive_public_key_bytes())); | ||
| // Edwards cache will be computed lazily on first signature | ||
| let private_key = PrivateKey::from(PrivateKeyData::DjbPrivateKey { | ||
| key, | ||
| edwards_cache: OnceLock::new(), | ||
| }); |
There was a problem hiding this comment.
Lazy keygen isn’t actually lazy with current new_without_cache
KeyPair::generate assumes curve25519::PrivateKey::new_without_cache avoids the Edwards computation, but the current implementation (curve25519.rs lines 55–71) still calls compute_ed_public_key, so the expensive scalar multiplication still happens and the comment here is misleading. This defeats the “lazy init” goal for key generation.
Consider updating new_without_cache to skip Edwards computation and use dummy values (similar to from_bytes_without_cache), or rename/adjust the API so it’s not advertised as lazy.
🔧 Proposed fix in wacore/libsignal/src/core/curve25519.rs
pub fn new_without_cache<R>(csprng: &mut R) -> Self
where
R: CryptoRng + Rng,
{
let mut bytes = [0u8; 32];
csprng.fill_bytes(&mut bytes);
bytes = scalar::clamp_integer(bytes);
let secret = StaticSecret::from(bytes);
- let (ed_public_key, sign_bit) = Self::compute_ed_public_key(&secret);
- PrivateKey {
- secret,
- ed_public_key,
- sign_bit,
- }
+ // Avoid Edwards computation here; compute lazily on first signature.
+ PrivateKey {
+ secret,
+ ed_public_key: CompressedEdwardsY::default(),
+ sign_bit: 0,
+ }
}🤖 Prompt for AI Agents
In `@wacore/libsignal/src/core/curve.rs` around lines 366 - 378, The comment is
correct: KeyPair::generate currently calls
curve25519::PrivateKey::new_without_cache expecting no Edwards computation, but
new_without_cache still calls compute_ed_public_key; update the implementation
so it truly defers expensive scalar multiplication (or change the call site).
Specifically, modify curve25519::PrivateKey::new_without_cache to mirror
from_bytes_without_cache by NOT calling compute_ed_public_key and instead
initializing the edwards cache as an empty OnceLock/dummy placeholder, or
alternatively change KeyPair::generate to call from_bytes_without_cache (or
another clearly-named factory) so the API behavior matches the “lazy” claim
(referencing PrivateKey::new_without_cache,
PrivateKey::from_bytes_without_cache, and KeyPair::generate).
This was referenced Jan 24, 2026
Merged
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR optimizes XEdDSA signature creation by caching the Edwards public key point, avoiding an expensive scalar multiplication on every signature call.
As noted in the XEdDSA specification:
Changes
XEDDSA_HASH_PREFIXstatic constant (minor optimization)ed_public_key(CompressedEdwardsY) andsign_bitusing lazy initialization viaOnceLockfrom_bytes_with_cache()andnew_without_cache()constructors for flexible key creationBencher CI Benchmark Results
Measured with
iai-callgrindin isolated CI environment (instruction count):Key Improvements
Key Generation (Lazy Initialization)
Real-World Impact by Use Case
Design: Lazy Initialization
Unlike the original approach that eagerly computed the Edwards cache at key creation (causing 2x key generation overhead), this implementation uses
OnceLockfor lazy initialization:This gives the best of both worlds - fast key generation AND fast signing.
Further Optimization Opportunities
Analysis of the libsignal crate revealed additional optimization opportunities for future PRs:
High Priority
crypto/hash.rsfinalize()always allocates Vec even when fixed-size array sufficesprotocol/state/session.rs:238protocol/state/session.rs:254Chainstructs inget_receiver_chain()Already Well-Optimized ✓
session_cipher.rs,group_cipher.rscrypto/aes_cbc.rsprotocol/ratchet/keys.rsfinalize_reset()to avoid recreating HMACprotocol/ratchet/keys.rsMessageKeyGeneratorenumcurve25519.rsvartime_double_scalar_mul_basepoint+ct_eqTest plan
cargo test --all)cargo clippy --all-targets)Summary by CodeRabbit
New Features
Performance Improvements
Tests
Chores
✏️ Tip: You can customize this high-level summary in your review settings.